import polars as pl
import time
from datetime import timedelta
import os

# pl.Config.set_global_thread_pool_size(8)  # 只使用8线程

def load_csv_fast(csv_path: str):
    """使用 Polars 多线程高效加载超大 CSV 文件"""
    if not os.path.exists(csv_path):
        raise FileNotFoundError(f"❌ 文件不存在：{csv_path}")

    print(f"📂 正在加载文件: {csv_path}")
    start = time.time()

    # Polars 的 scan_csv() 支持懒加载（可视为“流式 + 并行”）
    df_lazy = pl.scan_csv(
        csv_path,
        has_header=False,           # 没有表头
        new_columns=["user_id", "item_id", "category_id", "behavior_type", "timestamp"],
        infer_schema_length=1000,   # 采样1000行推断类型（足够快）
        ignore_errors=True,
        low_memory=True
    )

    # collect() 会触发真正的加载（自动使用多线程）
    df = df_lazy.collect()

    end = time.time()
    print(f"✅ 加载完成！总行数: {df.height:,}")
    print(f"⏱ 总耗时: {timedelta(seconds=end - start)}")

    return df


if __name__ == "__main__":
    csv_path = r"D:\data\公开数据集\淘宝用户行为数据\UserBehavior.csv"
    df = load_csv_fast(csv_path)

    print("\n📊 数据预览：")
    print(df.head(10))
    print("\n📈 各列类型：")
    print(df.schema)
    # print(100_000_000)  # 测试下大数字写法
